package com.lin.sqltool.service.impl;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.lin.sqltool.config.AISystemPrompt;
import com.lin.sqltool.config.ZhiPuAiConfig;
import com.lin.sqltool.factory.ConnectionServiceFactory;
import com.lin.sqltool.model.ConnectionInfoModel;
import com.lin.sqltool.service.AIService;
import com.lin.sqltool.service.ConnectionService;
import com.lin.sqltool.utils.CommonUtil;
import com.lin.sqltool.utils.DatabaseInfoSimplifyUtil;
import com.lin.sqltool.utils.MarkdownUtils;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.ChatCompletionRequest;
import com.zhipu.oapi.service.v4.model.ChatMessage;
import com.zhipu.oapi.service.v4.model.ChatMessageRole;
import com.zhipu.oapi.service.v4.model.ModelApiResponse;
import javafx.scene.control.Alert;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class ZhiPuAIServiceImpl implements AIService {


    /**
     * 用户上下文
     */
    private final Map<String, List<ChatMessage>> zhiPuAiChatMessages = new HashMap<String, List<ChatMessage>>();


    private final Map<String, String> databaseInfoCache = new HashMap<>();

    private static ClientV4 getClientV4() {
        return new ClientV4.Builder(ZhiPuAiConfig.API_KEY).build();
    }

    /**
     * 生成SQL
     *
     * @param requestId 请求的唯一id
     * @param info      连接信息
     * @param database  用户选中的数据库
     * @param prompt    文本
     * @return sql
     */
    @Override
    public String generateSQL(String requestId, ConnectionInfoModel info, String database, String prompt) {

        ConnectionService service =  ConnectionServiceFactory.getService(info.getDatabaseType());
        String dbInfo = databaseInfoCache.get(requestId);
        if (dbInfo == null) {
            try {
                dbInfo = DatabaseInfoSimplifyUtil.simplify(service.getDatabaseInfo(info, database));
            } catch (JsonProcessingException e) {
                CommonUtil.showAlert(Alert.AlertType.ERROR, "数据库信息查询错误", e.getMessage());
            }
        }

        List<ChatMessage> messages = zhiPuAiChatMessages.get(requestId);    //先看缓存有没有聊天记录
        if (messages == null) {
            messages = new ArrayList<>();
            messages.add(new ChatMessage(ChatMessageRole.SYSTEM.value(), AISystemPrompt.generateSQL + dbInfo)); //如果没有重新组成聊天记录
        }
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), prompt);
        messages.add(chatMessage);
        System.out.println("requestId: " + requestId);
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(ZhiPuAiConfig.MODEL_NAME)
                .stream(Boolean.FALSE)
                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .requestId(requestId)
                .build();
        ModelApiResponse invokeModelApiResp = getClientV4().invokeModelApi(chatCompletionRequest);
        String result = invokeModelApiResp.getData().getChoices().get(0).getMessage().getContent().toString();
        return MarkdownUtils.removeMarkdownTags(result);
    }


}
